#include <iostream>
#include <vector>

using namespace std;

class Node
{
public:
    int val;
    Node *left;
    Node *right;
    Node *next;

    Node() : val(0), left(NULL), right(NULL), next(NULL) {}

    Node(int _val) : val(_val), left(NULL), right(NULL), next(NULL) {}

    Node(int _val, Node *_left, Node *_right)
        : val(_val), left(_left), right(_right), next(NULL) {}
};

// TODO: can use next as a link list, to remvoe levels[0]

class Solution
{
public:
    Node *connect(Node *root)
    {
        if (root == nullptr)
            return nullptr;
        vector<vector<Node *>> levels(2);
        levels[0].push_back(root);
        while (levels[0].size() != 0)
        {
            for (auto n : levels[0])
            {
                if (n->left != nullptr)
                    levels[1].push_back(n->left);
                if (n->right != nullptr)
                    levels[1].push_back(n->right);
            }
            if (levels[1].empty())
            {
                break;
            }
            for (int i = 0; i < levels[1].size() - 1; i++)
            {
                levels[1][i]->next = levels[1][i + 1];
            }
            levels[0] = levels[1];
            levels[1].clear();
        }
        return root;
    }
};

int main()
{
    cout << "Code117" << endl;
    Solution s;
    Node root(1, new Node(2, new Node(4), new Node(5)), new Node(3, nullptr, new Node(7)));
    Node* ret = s.connect(&root);
    return 0;
}